import numpy as np
import time


class KKT2:

    def __init__(self, env):
        self.env = env
        self.action, self.bandwidth, self.frequency = None, None, None

    def run(self, collector):
        for i in range(200):
            print("**************************************************************************")
            print(f"Current episode: {i}")
            start_timestamp = time.time()
            # 1. 生成卸载决策（a）
            self.choose_action()
            # 2. 进行资源分配（B, f）
            self.bandwidth, self.frequency = self.ra(self.action)
            # 3. 计算当前决策的总成本
            self.env.execute(action=self.action, bandwidth=self.bandwidth, frequency=self.frequency)
            cost = self.env.calc_cost()
            t_ratio, cpu_ratio = self.env.calc_utilization()
            collector.collect(cost, sum(self.action), t_ratio, cpu_ratio)
            collector.next()
            print(f"Offloading cost: {cost}, Offloading num: {sum(self.action)}, Time Ratio:{t_ratio}, CPU Ratio:{cpu_ratio}, Time:{round(time.time() - start_timestamp, 2)}")
            # 4. 计算资源利用率
            ratio_b = 0
            ratio_f = 0
            for g in self.env.parts_ues:
                ratio_b += sum(self.bandwidth[g]) / self.env.config.uav_bandwidth
                ratio_f += sum(self.frequency[g]) / self.env.config.uav_cpu_frequency
            print(f"Avg resource 1: {ratio_b / self.env.G}, 2: {ratio_f / self.env.G}")
        return self.action, self.bandwidth, self.frequency

    def choose_action(self):
        if self.action is None:
            self.action = np.zeros(self.env.M, dtype=int)
        else:
            cost = 9999
            action_opt = self.action
            for i in range(self.env.M):
                action = self.action.copy()
                action[i] = (0 if action[i] else 1)
                bandwidth, frequency = self.ra_without_constraint(action=action)
                ue_p = self.check_timeout(action, bandwidth, frequency)
                while len(ue_p) != 0:
                    ue_p_parts = self.group_part(ue_p)
                    for part in ue_p_parts:
                        part_ues_timeout = ue_p_parts[part]
                        max_t_loc = np.argmax(self.env.meta_ue_loc[part_ues_timeout])
                        action[part_ues_timeout[max_t_loc]] = 0
                    bandwidth, frequency = self.ra_without_constraint(action=action)
                    ue_p = self.check_timeout(action, bandwidth, frequency)
                c = self.env.calc_cost(action=action, bandwidth=bandwidth, frequency=frequency)
                if c < cost:
                    action_opt = action
                    cost = c
            self.action = action_opt

    def ra(self, action):
        # 1. Calculate timeout UEs
        bandwidth, frequency = self.ra_without_constraint(action)
        ue_p = self.check_timeout(action, bandwidth, frequency)
        if len(ue_p) == 0:
            return bandwidth, frequency
        ue_p_parts = self.group_part(ue_p)
        for part in ue_p_parts:
            part_ue_p = ue_p_parts[part]
            part_ue_q = [i for i in self.env.parts_ues[part] if action[i] and i not in part_ue_p]
            # 2. Calculate nu1, nu2
            k1 = self.env.meta_trans_rate / self.env.meta_ue_loc
            k2 = self.env.meta_task_cpu / self.env.meta_ue_loc
            nu2_1 = self.env.config.uav_bandwidth * self.env.config.uav_hover_time - sum(self.env.meta_trans_rate[part_ue_p])
            # nu2_2 = sum(np.sqrt(k1[part_ue_q])) / nu2_1
            nu2_5 = sum(np.sqrt(self.env.meta_trans_rate[part_ue_p] * self.env.meta_task_cpu[part_ue_p]))
            nu2_2 = self.env.config.uav_hover_time * sum(np.sqrt(k1[part_ue_q])) * nu2_5 / nu2_1
            nu2_3 = self.env.config.uav_hover_time * sum(np.sqrt(k2[part_ue_q]))
            nu2_4 = self.env.config.uav_cpu_frequency * self.env.config.uav_hover_time - sum(self.env.meta_task_cpu[part_ue_p])
            nu2_6 = (nu2_5 ** 2) / nu2_1
            nu2 = (nu2_2 + nu2_3) / (nu2_4 - nu2_6)
            nu1 = (np.sqrt(nu2) * nu2_5 + self.env.config.uav_hover_time * sum(np.sqrt(k1[part_ue_q]))) / nu2_1
            nu1 = nu1 ** 2
            # 3. Calculate lambda, bandwidth, frequency
            lambdas = np.zeros(self.env.M, dtype=int)
            for idx in part_ue_p:
                lambdas[idx] = (np.sqrt(nu1 * k1[idx]) + np.sqrt(nu2 * k2[idx])) * self.env.meta_ue_loc[idx] / self.env.config.uav_hover_time
                lambdas[idx] = lambdas[idx] ** 2 - 1
                bandwidth[idx] = np.sqrt(k1[idx] * (1 + lambdas[idx]) / nu1)
                frequency[idx] = np.sqrt(k2[idx] * (1 + lambdas[idx]) / nu2)
            for idx in part_ue_q:
                bandwidth[idx] = np.sqrt(k1[idx] / nu1)
                frequency[idx] = np.sqrt(k2[idx] / nu2)
        return bandwidth, frequency

    def group_part(self, ue_p):
        ue_p_parts = {}
        for idx in ue_p:
            part = self.env.ues[idx].part
            if part in ue_p_parts:
                ue_p_parts[part].append(idx)
            else:
                ue_p_parts[part] = [idx]
        return ue_p_parts

    def ra_without_constraint(self, action):
        bandwidth = np.zeros(self.env.M, dtype=int)
        frequency = np.zeros(self.env.M, dtype=int)
        for g in range(self.env.G):
            off_idx = [idx for idx in self.env.parts_ues[g] if action[idx]]
            meta_trans_rate = np.sqrt(self.env.meta_trans_rate)
            meta_task_cpu = np.sqrt(self.env.meta_task_cpu)
            for ue_idx in off_idx:
                bandwidth[ue_idx] = meta_trans_rate[ue_idx] * self.env.config.uav_bandwidth / sum(meta_trans_rate[off_idx])
                frequency[ue_idx] = meta_task_cpu[ue_idx] * self.env.config.uav_cpu_frequency / sum(meta_task_cpu[off_idx])
        return bandwidth, frequency

    def check_timeout(self, action, bandwidth, frequency, ues_idx=None):
        if ues_idx is None:
            ues_idx = [x for x in range(self.env.M)]
        with np.errstate(divide='ignore'):
            t_off = self.env.meta_trans_rate[ues_idx] / bandwidth
            t_exe = self.env.meta_task_cpu[ues_idx] / frequency
        for m in range(len(action)):
            if action[m] == 0:
                t_off[m] = 0
                t_exe[m] = 0
        t_serv = t_off + t_exe
        timeout_result = (t_serv > self.env.config.uav_hover_time)
        # noinspection PyTypeChecker
        return [idx for idx, val in enumerate(timeout_result) if val]
